

import os
import random
from tqdm import tqdm
import shutil
import os.path
jpg_path = r'/home/xys/datasets/Exdark/ExDark-all/img/'  #jpg和txt混合路径
xml_path = r'/home/xys/datasets/Exdark/ExDark-all/ANN/'     #txt保存路径
save_jpg_path = r'/home/xys/datasets/dataset_che/dataset-coco/dataset-ExDark/train/'       #图片路径
save_xml_path = r'/home/xys/datasets/dataset_che/dataset-coco/dataset-ExDark/train-xml/'
save_jpg_path_val = r'/home/xys/datasets/dataset_che/dataset-coco/dataset-ExDark/val/'
save_xml_path_val = r'/home/xys/datasets/dataset_che/dataset-coco/dataset-ExDark/val-xml/'
save_jpg_path_test = r'/home/xys/datasets/dataset_che/dataset-coco/dataset-ExDark/test/'
save_xml_path_test = r'/home/xys/datasets/dataset_che/dataset-coco/dataset-ExDark/test-xml/'

txt_path = r'/home/xys/datasets/Exdark/ExDark-all/imageclasslist.txt'

# file_names = os.listdir(jpg_path)
# lens =  len(file_names)
clas = []
with open(txt_path, "r") as f:
    next(f)        #跳一行
    for line in f.readlines():
        line = line.strip('\n')
        list = line.split(" ")
        # print(list)
        clas.append(list)

for i in range(len(clas)):
    if int(clas[i][4]) == 1:
        shutil.copy(jpg_path+clas[i][0],save_jpg_path)
        # os.rename(save_jpg_path+clas[i][0],save_jpg_path+clas[i][0].split('.')[:-1]+'.jpg')
        x = clas[i][0].split('.')[0]
        shutil.copy(os.path.join(xml_path+x+'.xml'),save_xml_path)

    elif int(clas[i][4]) == 2:
        shutil.copy(jpg_path+clas[i][0],save_jpg_path_val)
        y = clas[i][0].split('.')[0]
        # os.rename(save_jpg_path + clas[i][0], save_jpg_path + clas[i][0].split('.')[:-1] + '.jpg')
        shutil.copy(os.path.join(xml_path+y+'.xml'),save_xml_path_val)

    elif int(clas[i][4]) == 3:
        shutil.copy(jpg_path+clas[i][0],save_jpg_path_test)
        z = clas[i][0].split('.')[0]
        # os.rename(save_jpg_path + clas[i][0], save_jpg_path + clas[i][0].split('.')[:-1] + '.jpg')
        shutil.copy(os.path.join(xml_path+z+'.xml'),save_xml_path_test)













# random.shuffle(file_names)
# lens = len(file_names)
# for file in tqdm(file_names[0:int(0.8*lens)]):
#     file_back = file[-4:]
#     file_front = file[:-4]
#     if file_back == ".jpg":
#         # os.remove(mix_file_path + file)
#         shutil.copy(os.path.join(jpg_path, file), save_jpg_path)
#         shutil.copy(os.path.join(xml_path,file_front+'.xml'),save_xml_path)
# for file in tqdm(file_names[int(0.8*lens):int(0.9*lens)]):
#     file_back = file[-4:]
#     file_front = file[:-4]
#     if file_back == ".jpg":
#         # os.remove(mix_file_path + file)
#         shutil.copy(os.path.join(jpg_path, file), save_jpg_path_val)
#         shutil.copy(os.path.join(xml_path,file_front+'.xml'),save_xml_path_val)
# for file in tqdm(file_names[int(0.9*lens):lens]):
#     file_back = file[-4:]
#     file_front = file[:-4]
#     if file_back == ".jpg":
#         # os.remove(mix_file_path + file)
#         shutil.copy(os.path.join(jpg_path, file), save_jpg_path_test)
#         shutil.copy(os.path.join(xml_path,file_front+'.xml'),save_xml_path_test)